/* Copyright (c) 2024 Huawei Technologies Co., Ltd.
* This file is a part of the CANN Open Software.
* Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
* ===================================================================================================================*/

#include <functional>
#include "graph/utils/args_format_desc_utils.h"
#include "common/checker.h"
#include "common/ge_common/debug/ge_log.h"
#include "graph/args_format_desc.h"

namespace {
constexpr int32_t kDecimalCarry = 10;
constexpr int32_t kAsciiZero = 48;

using ParseFunc = std::function<ge::graphStatus(const ge::OpDescPtr &, const std::string &, const ge::AddrType type,
                                            std::vector<ge::ArgDesc> &)>;
using SerializeFunc =
    std::function<void(std::stringstream &ss, const std::string &pattern, const ge::ArgDesc &arg_desc)>;

struct PatternHandler {
  ge::AddrType type;
  ParseFunc parse;
  SerializeFunc serialize;
};

const std::map<ge::AddrType, std::string> kAddrType2String = {
    {ge::AddrType::INPUT, "i"},
    {ge::AddrType::OUTPUT, "o"},
    {ge::AddrType::WORKSPACE, "ws"},
    {ge::AddrType::TILING, "t"},
    {ge::AddrType::INPUT_DESC, "i_desc"},
    {ge::AddrType::OUTPUT_DESC, "o_desc"},
    {ge::AddrType::FFTS_ADDR, "ffts_addr"},
    {ge::AddrType::OVERFLOW_ADDR, "overflow_addr"},
    {ge::AddrType::TILING_FFTS, "t_ffts"},
    {ge::AddrType::HIDDEN_INPUT, "hi"},
    {ge::AddrType::OP_TYPE, "*op_type"},
    {ge::AddrType::TILING_CONTEXT, "tiling_context"},
    {ge::AddrType::PLACEHOLDER, ""},
};

// 处理输入、输出、workspace
ge::graphStatus MultiAddrsParser(const ge::OpDescPtr &op_desc, const std::string &pattern_str, const ge::AddrType type,
                                 std::vector<ge::ArgDesc> &arg_descs) {
  (void) op_desc;
  auto iter = kAddrType2String.find(type);
  // 处理特殊场景，即: i*, o*, ws*场景，即存在folded场景
  if ((iter != kAddrType2String.end()) && (iter->second + "*" == pattern_str)) {
    arg_descs.push_back({type, -1, false, {0U}});
  } else {  // 处理：i0, o0，ws0， i0*, o0*等场景
    int32_t ir_idx{0};
    bool has_idx{false};
    for (size_t i = iter->second.size(); i < pattern_str.size(); ++i) {
      if (isdigit(pattern_str[i])) {
        ir_idx = ir_idx * kDecimalCarry + static_cast<int32_t>(pattern_str[i]) - kAsciiZero;
        has_idx = true;
      }
    }
    GE_ASSERT(has_idx, "Arg format [%s] is invalid", pattern_str.c_str());

    bool folded = pattern_str[pattern_str.length() - 1UL] != '*';
    arg_descs.push_back({type, ir_idx, folded, {0U}});
  }
  return ge::GRAPH_SUCCESS;
}

struct PatternCmp {
  bool operator()(const std::string &lhs, const std::string &rhs) const {
    return lhs.size() >= rhs.size();
  };
};
// 会按照key值长短排序，保证匹配优先匹配长字符串序列
static const std::map<std::string, PatternHandler, PatternCmp> kPatternToHandler = {
    {"i", {ge::AddrType::INPUT, MultiAddrsParser, ge::ArgsFormatDesc::ArrayLikeSerializer}},
    {"o", {ge::AddrType::OUTPUT, MultiAddrsParser, ge::ArgsFormatDesc::ArrayLikeSerializer}},
    {"ws", {ge::AddrType::WORKSPACE, MultiAddrsParser, ge::ArgsFormatDesc::ArrayLikeSerializer}},
    {"t", {ge::AddrType::TILING, ge::ArgsFormatDesc::DefaultParser, ge::ArgsFormatDesc::DefaultSerializer}},
    {"i_desc", {ge::AddrType::INPUT_DESC, ge::ArgsFormatDesc::IODescParser, ge::ArgsFormatDesc::ArrayLikeSerializer}},
    {"o_desc", {ge::AddrType::OUTPUT_DESC, ge::ArgsFormatDesc::IODescParser, ge::ArgsFormatDesc::ArrayLikeSerializer}},
    {"ffts_addr", {ge::AddrType::FFTS_ADDR, ge::ArgsFormatDesc::DefaultParser, ge::ArgsFormatDesc::DefaultSerializer}},
    {"overflow_addr",
     {ge::AddrType::OVERFLOW_ADDR, ge::ArgsFormatDesc::DefaultParser, ge::ArgsFormatDesc::DefaultSerializer}},
    {"t_ffts",
     {ge::AddrType::TILING_FFTS, ge::ArgsFormatDesc::DefaultParser, ge::ArgsFormatDesc::FftsTilingSerializer}},
    {"hi",
     {ge::AddrType::HIDDEN_INPUT, ge::ArgsFormatDesc::HiddenInputParser, ge::ArgsFormatDesc::HiddenInputSerializer}},
    {"*op_type", {ge::AddrType::OP_TYPE, ge::ArgsFormatDesc::DefaultParser, ge::ArgsFormatDesc::DefaultSerializer}},
    {"tiling_context",
     {ge::AddrType::TILING_CONTEXT, ge::ArgsFormatDesc::TilingContextParser,
      ge::ArgsFormatDesc::TilingContextSerializer}},
    {"", {ge::AddrType::PLACEHOLDER, ge::ArgsFormatDesc::PlaceholderParser, ge::ArgsFormatDesc::DefaultSerializer}},
};

}  // namespace

namespace ge {

void ArgsFormatDescUtils::Append(std::vector<ArgDesc> &arg_descs, ge::AddrType type, int32_t ir_idx, bool folded) {
  arg_descs.push_back({type, ir_idx, folded, {0}});
}

void ArgsFormatDescUtils::AppendHiddenInput(std::vector<ArgDesc> &arg_descs, ge::HiddenInputType hidden_type) {
  arg_descs.push_back({AddrType::HIDDEN_INPUT, static_cast<int32_t>(hidden_type), false, {0}});
}

void ArgsFormatDescUtils::AppendTilingContext(std::vector<ArgDesc> &arg_descs, ge::TilingContextSubType sub_type) {
  arg_descs.push_back({AddrType::TILING_CONTEXT, static_cast<int32_t>(sub_type), false, {0}});
}

graphStatus ArgsFormatDescUtils::Parse(const std::string &str, std::vector<ArgDesc> &arg_descs) {
  arg_descs.clear();
  size_t start_idx = 0UL;
  while (start_idx < str.size()) {
    GE_ASSERT(str[start_idx] == '{', "SyntaxError: argsformat should be surrounded by '{','}'");
    size_t end_idx = start_idx + 1UL;
    bool parsed{false};
    while (end_idx < str.size()) {
      if (str[end_idx] == '}') {
        std::string pattern_str = str.substr(start_idx + 1, end_idx - start_idx - 1);
        for (const auto &iter : kPatternToHandler) {
          if (strncmp(pattern_str.c_str(), iter.first.c_str(), iter.first.length()) == 0) {
            GE_ASSERT_SUCCESS(iter.second.parse(nullptr, pattern_str, iter.second.type, arg_descs),
                              "args format [%s] parser failed.", pattern_str.c_str());
            parsed = true;
            start_idx = end_idx + 1UL;
            break;
          }
        }
      }
      ++end_idx;
    }
    GE_ASSERT(parsed, "SyntaxError: argsformat should be surrounded by '{','}'");
  }
  return GRAPH_SUCCESS;
}

std::string ArgsFormatDescUtils::Serialize(const std::vector<ArgDesc> &arg_descs) {
  std::stringstream ss;
  for (const auto &arg_desc : arg_descs) {
    for (const auto &iter : kPatternToHandler) {
      if (iter.second.type == arg_desc.addr_type) {
        ss << '{';
        iter.second.serialize(ss, iter.first, arg_desc);
        ss << '}';
      }
    }
  }
  return ss.str();
}

std::string ArgsFormatDescUtils::ToString(const std::vector<ArgDesc> &arg_descs) {
  return Serialize(arg_descs);
}
}  // namespace ge